from keras.models import load_model

from utils.boxes import create_prior_boxes
from utils.boxes import to_point_form
from utils.inference import detect, detect_self
from utils.training import MultiboxLoss


from utils.data_management import get_class_names

import cv2 as cv 
import numpy as np
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

def draw_box(img, pt1, pt2, scores):
    if scores >= thr:
        string = '%f'%(scores)
        cv.rectangle(img, pt1, pt2, color=(0, 255, 0))
        cv.putText(img, string, (int(pt1[0]), int(pt1[1])), cv.FONT_HERSHEY_COMPLEX, 1, (0,255,0), 1)
    cv.imshow('img', img)
    cv.waitKey(5)

def predict_box(raw):
    img = raw.astype('float32') / 128.0 - 1

    img = np.expand_dims(img, axis=0)
    pred = model.predict(img)[0]
    w = raw.shape[-2]
    h = raw.shape[-3]
    all_boxes = []
    detections = detect_self(pred, prior_boxes)
    for j in range(len(pred[:,4:][0])):
        dets = detections[0, j, :]
        mask = np.squeeze(dets[:,0] > 0.01)
        dets = dets[mask]
        if len(dets) == 0:
            continue

        boxes = dets[:, 1:]
        boxes[:, 0] *= w 
        boxes[:, 2] *= w
        boxes[:, 1] *= h 
        boxes[:, 3] *= h 

        scores = dets[:, 0]
        label = np.reshape(np.array([j]*len(boxes)),(len(boxes),1))
        cls_dets = np.hstack((label, boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
        all_boxes.append(cls_dets) 
    return all_boxes, raw

def predict_and_draw(img):
    img = cv.resize(img, (input_w, input_h))
    all_boxes, raw = predict_box(img)
    print("Print top %d results!"%top_k) 
    for i, boxes in enumerate(all_boxes[1:]):
        temp = raw.copy().squeeze()
        for points in boxes[:top_k]:
            label = points[0]
            pt1 = tuple(points[1:3])
            pt2 = tuple(points[3:5])
            scores = points[-1]
            draw_box(temp, pt1, pt2, scores)

if __name__ == '__main__':
    from model import main_model
    from_files = False

    multibox_loss = MultiboxLoss(2, 3, 0.1) 
    model = load_model(r'C:\Users\nxf48054\Desktop\person_detect/face_detect_ssd/trained_models/ssd_face_detect/weights.46-2.68.h5', custom_objects={'compute_loss':multibox_loss.compute_loss})
    # model.save_weights("./weights.26-2.33.hdf5")
    #model.summary()
    img_path = r"F:\Face-detection-with-mobilenet-ssd-master\dataset/out_train_multi/"
    _, input_w, input_h, _ = model.input.shape.as_list()
  
    # val model
    # val_model = main_model((192, 256, 3), 2, 'Val')[0]
    # val_model.load_weights("./weights.26-2.33.hdf5", by_name=True)
    # val_model.save("./trained_models/weights.26-2.33_post_process.h5")

    # our code need point format
    from utils.boxes import create_prior_boxes
    from train import get_configuration_file
    prior_boxes = (create_prior_boxes(configuration=get_configuration_file()))
    np.save("prior_boxes", prior_boxes)
    prior_boxes = to_point_form(prior_boxes)
    # the max box 
    top_k = 100  
    thr = 0.6
    empty_buffer = np.empty((25, 256, 256, 3), dtype='uint8')
    idx = 0

    cv.namedWindow("img", 0)

    if from_files:
        for p in os.listdir(img_path)[1:]:
            img = cv.imread(os.path.join(img_path, p))
            # empty_buffer[idx] = img
            # idx += 1
            # if(idx == 25):
            #     break
            # cv.imwrite("1.jpg", img)
            predict_and_draw(img)
    else:
        capture = cv.VideoCapture(0)
        while(1):
            ret, frame = capture.read()
            predict_and_draw(frame) 
    
    buffer_b = empty_buffer.tobytes()
    with open("25_people.bin", "wb") as f:
        f.write(buffer_b)
        f.close()
    print('done!')



